"""
1. 获取3大指数的日线数据
2. 对3个指数分别计算N日涨幅，取涨幅最大的那个指数为目标仓位
3. 检查已有持仓，如果不一样就调整
4. 每日执行第4，5步。
"""
from jqdatasdk import *
import pandas as pd
import numpy as np
import os
from modular.context import Context
from modular.strategy import Strategy
from modular.operation import Operation
from modular.visualization import Visualization

username = "13021491353"
password = "AInvest7937"
auth(username=username,password=password)

def backtest(delay=20):
    init_money = 10000
    my_delay = delay
    start_date = "2017-12-01"
    end_date = "2020-12-31"
    # index_ls 元素的顺序必须和 index_etf_ls 的元素顺序一样
    use_index = [0,1,2] # 使用哪几个指数计算
    index_ls = ['000300.XSHG',#沪深300
                '000905.XSHG',#中证500
                '000016.XSHG',#上证50
                ]
    index_etf_ls = ['510300.XSHG', # 沪深300ETF
                '510500.XSHG', # 中证500ETF
                '510050.XSHG', # 上证50ETF
                ]
    data_dir = r'./data/'
    index_filename = r'3main_index_data_'+start_date[:4]+'-'+end_date[:4]+'.csv'
    index_etf_filename = r'3main_index_etf_data_'+start_date[:4]+'-'+end_date[:4]+'.csv'

    # 获取3大指数及对应ETF基金的数据
    if (    os.path.exists(os.path.join(data_dir,index_filename)) and
            os.path.exists(os.path.join(data_dir,index_etf_filename)) ):
        index_data_df = pd.read_csv(os.path.join(data_dir,index_filename),index_col=0)
        index_etf_data_df = pd.read_csv(os.path.join(data_dir,index_etf_filename),index_col=0)

    else:
        index_data_df = pd.DataFrame()
        index_etf_data_df = pd.DataFrame()
        for i in range(len(index_ls)):
            # 获取指数数据,以df形式保存
            i_data = get_price(security=index_ls[i], start_date=start_date, end_date=end_date, fields=["close"],
                               frequency="1d")
            index_data_df = pd.concat([index_data_df, i_data], axis=1)
            # 获取etf数据,以df形式保存
            i_etf_data = get_price(security=index_etf_ls[i], start_date=start_date, end_date=end_date, fields=["close"],
                                   frequency="1d")
            index_etf_data_df = pd.concat([index_etf_data_df, i_etf_data], axis=1)
        index_data_df.columns = index_ls
        index_data_df.to_csv(os.path.join(data_dir, index_filename))
        index_etf_data_df.columns = index_etf_ls
        index_etf_data_df.to_csv(os.path.join(data_dir, index_etf_filename))
        index_data_df = pd.read_csv(os.path.join(data_dir,index_filename),index_col=0)
        index_etf_data_df = pd.read_csv(os.path.join(data_dir,index_etf_filename),index_col=0)
    index_data = {}
    index_etf_data = {}
    for i_index in index_ls:
        index_data[i_index] = index_data_df[i_index]
    for i_etf in index_etf_ls:
        index_etf_data[i_etf] = index_etf_data_df[i_etf]

    date_range = index_data[index_ls[0]].index

    my_context = Context(currency=init_money)
    my_strategy = Strategy()
    my_opt = Operation()

    # 每天计算目标仓位，并进行调整
    my_total_ls = []
    for i_dt in date_range:
        # 调整仓位
        my_opt.update_account(context=my_context, invest_data=index_etf_data, dt=i_dt)
        my_opt.adjust_have(context=my_context,dt=i_dt,invest_data=index_etf_data)
        my_total_ls.append(my_context.total)
        # 重置目标仓位
        my_context.target_have = []
        # 每一天对三个指数进行计算
        judge_result = my_strategy.judge2(index_data_df=index_data_df.iloc[:,use_index],dt=i_dt,delay=my_delay)
        if judge_result == None:
            # 没有N天前的数据，或者当天应该空仓
            continue
        else: # 有要买的指数
            buy_index = index_etf_ls[index_ls.index(judge_result)]
            my_context.target_have.append(buy_index)
        # 止盈止损
        if len(my_context.now_have)!=0:
            if len(my_context.now_have)!=1:
                print('dt : '+str(i_dt))
                print(my_context.now_have.items())
            else:
                is_stop_win = my_strategy.is_stop_win(context=my_context,
                                                      invest_code=list(my_context.now_have.keys())[0],
                                                      index_etf_data_df=index_etf_data_df,
                                                      dt=i_dt,
                                                      stop_win_rate=0.05)
                is_stop_loss = my_strategy.is_stop_loss(context=my_context,
                                                        invest_code=list(my_context.now_have.keys())[0],
                                                        index_etf_data_df=index_etf_data_df,
                                                        dt=i_dt,
                                                        stop_loss_rate=0.05)
                if (is_stop_loss==True or is_stop_win==True):
                    if list(my_context.now_have.keys())[0] in my_context.target_have:
                        print('---------------------------')
                        print(str(i_dt)+' trigger stop loss/win')
                        print(my_context.target_have)
                        my_context.target_have.remove(list(my_context.now_have.keys())[0])

    # 结果可视化
    my_visual = Visualization()
    my_total_ls = np.array(my_total_ls)/my_total_ls[0] # 归一化
    my_visual.plot(x=date_range,y=my_total_ls,etf_df=index_etf_data_df,delay=my_delay,use_index=use_index)
    print('==================')

if __name__=='__main__':
    for dl in range(1,30):
        backtest(delay=dl)